import cv2
import numpy as np

class LSB:
    def __init__(self):
        # 使用二进制结束标记，避免字符编码问题
        self.terminator_binary = '1111111111111110'  # 16位结束标记
    
    def embed(self, image_path, message, output_path):
        try:
            # 读取图像
            img = cv2.imread(image_path)
            if img is None:
                print(f"无法读取图像: {image_path}")
                return False
            
            print(f"图像形状: {img.shape}")
            print(f"原始消息: {message}")
            
            # 将消息转换为二进制（使用UTF-8编码避免中文问题）
            binary_message = ''.join([format(byte, '08b') for byte in message.encode('utf-8')])
            # 添加结束标记
            binary_message += self.terminator_binary
            
            print(f"二进制消息长度: {len(binary_message)} 位")
            
            # 检查图像容量
            total_pixels = img.shape[0] * img.shape[1] * 3
            if len(binary_message) > total_pixels:
                print(f"消息太长: {len(binary_message)} 位 > {total_pixels} 像素容量")
                return False
            
            print(f"图像容量: {total_pixels} 位，使用: {len(binary_message)} 位")
            
            # 嵌入消息
            data_index = 0
            height, width, channels = img.shape
            
            for i in range(height):
                for j in range(width):
                    for k in range(channels):  # BGR通道
                        if data_index < len(binary_message):
                            # 修改最低有效位
                            pixel_value = img[i, j, k]
                            new_pixel_value = (pixel_value & 0xFE) | int(binary_message[data_index])
                            img[i, j, k] = new_pixel_value
                            data_index += 1
                        else:
                            break
                    if data_index >= len(binary_message):
                        break
                if data_index >= len(binary_message):
                    break
            
            print(f"成功嵌入 {data_index} 位数据")
            
            # 保存为PNG格式（无损）
            if not output_path.lower().endswith('.png'):
                output_path = output_path.rsplit('.', 1)[0] + '.png'
                
            success = cv2.imwrite(output_path, img, [cv2.IMWRITE_PNG_COMPRESSION, 0])
            if not success:
                print(f"保存图像失败: {output_path}")
                return False
                
            print(f"图像保存成功: {output_path}")
            return True
            
        except Exception as e:
            print(f"嵌入过程中发生错误: {str(e)}")
            import traceback
            traceback.print_exc()
            return False
    
    def extract(self, image_path):
        try:
            # 读取图像
            img = cv2.imread(image_path)
            if img is None:
                print(f"无法读取图像: {image_path}")
                return ""
            
            print(f"提取图像形状: {img.shape}")
            
            # 提取二进制数据
            binary_data = ""
            height, width, channels = img.shape
            
            # 先提取所有LSB位
            for i in range(height):
                for j in range(width):
                    for k in range(channels):
                        lsb = str(img[i, j, k] & 1)
                        binary_data += lsb
            
            print(f"提取的二进制数据长度: {len(binary_data)} 位")
            
            # 查找结束标记
            terminator_pos = binary_data.find(self.terminator_binary)
            if terminator_pos == -1:
                print("未找到结束标记，尝试提取所有数据")
                # 如果没有找到结束标记，提取所有数据
                message_binary = binary_data
            else:
                print(f"找到结束标记在位置: {terminator_pos}")
                message_binary = binary_data[:terminator_pos]
            
            print(f"消息二进制长度: {len(message_binary)} 位")
            
            # 确保二进制长度是8的倍数
            if len(message_binary) % 8 != 0:
                message_binary = message_binary[:-(len(message_binary) % 8)]
                print(f"调整后二进制长度: {len(message_binary)} 位")
            
            # 转换为字节
            decoded_bytes = bytearray()
            for i in range(0, len(message_binary), 8):
                byte_str = message_binary[i:i+8]
                if len(byte_str) == 8:
                    byte_val = int(byte_str, 2)
                    decoded_bytes.append(byte_val)
            
            print(f"解码字节数: {len(decoded_bytes)}")
            
            # 尝试UTF-8解码
            try:
                decoded_message = decoded_bytes.decode('utf-8')
                print(f"成功解码消息: {decoded_message}")
                return decoded_message
            except UnicodeDecodeError as e:
                print(f"UTF-8解码失败: {e}")
                # 尝试其他编码或返回原始字节
                try:
                    decoded_message = decoded_bytes.decode('latin-1')
                    print(f"使用latin-1解码: {decoded_message}")
                    return decoded_message
                except:
                    # 返回十六进制表示
                    hex_representation = decoded_bytes.hex()
                    print(f"原始字节(十六进制): {hex_representation}")
                    return f"原始数据(十六进制): {hex_representation}"
            
        except Exception as e:
            print(f"提取过程中发生错误: {str(e)}")
            import traceback
            traceback.print_exc()
            return ""